package ai.koog.prompt.message

import kotlinx.datetime.Clock
import kotlinx.datetime.Instant
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.jsonObject
import kotlin.jvm.JvmOverloads

/**
 * Represents a message exchanged in a chat with LLM. Messages can be categorized
 * by their type and role, denoting the purpose and source of the message.
 *
 * Represents both a message from LLM and a message to LLM from user or environment.
 */
@Serializable
public sealed interface Message {
    /**
     * The content of the message.
     */
    public val content: String

    /**
     * The role associated with the message.
     */
    public val role: Role

    /**
     * Stores metadata information for the current message instance, such as token count and timestamp.
     */
    public val metaInfo: MessageMetaInfo

    /**
     * Represents a request message in the chat.
     */
    @Serializable
    public sealed interface Request : Message {
        override val metaInfo: RequestMetaInfo
    }

    /**
     * Represents a response message in the chat.
     */
    @Serializable
    public sealed interface Response : Message {
        override val metaInfo: ResponseMetaInfo

        /**
         * Creates a copy of the current Response instance with updated metadata.
         */
        public fun copy(updatedMetaInfo: ResponseMetaInfo): Response
    }

    /**
     * Defines the role of the message in the chat (e.g., system, user, assistant, tool).
     */
    @Serializable
    public enum class Role {
        /**
         * Role indicating a system message.
         */
        System,

        /**
         * Role for messages generated by the user.
         */
        User,

        /**
         * Role for messages generated by an assistant (e.g., an AI assistant).
         */
        Assistant,

        /**
         * Role for messages related to tools (e.g., tool usage or tool results).
         */
        Tool
    }

    /**
     * Represents a message that includes one or more attachments.
     *
     * This interface extends the [Message] interface, adding support for managing
     * additional content in the form of attachments. Attachments can be of various
     * types such as images, videos, audio files, or other file formats. Implementing
     * classes should define specific details of the message alongside its attachments.
     *
     * @property attachments A list of attachments associated with the message.
     * Each attachment contains specific content and metadata.
     */
    @Serializable
    public sealed interface WithAttachments : Message {
        /**
         * A list of attachments associated with a message.
         * Each attachment represents additional content that can be included,
         * such as images, videos, audio files, or other documents.
         */
        public val attachments: List<Attachment>
    }

    /**
     * Represents a message sent by the user as a request.
     *
     * @property content The main content of the user's message.
     * @property metaInfo Metadata associated with the request, including timestamp information. Defaults to a new [RequestMetaInfo].
     * @property attachments Optional media content attached to the message.
     * @property role The role of the message, which is fixed as [Role.User] for this implementation.
     */
    @Serializable
    public data class User @JvmOverloads constructor(
        override val content: String,
        override val metaInfo: RequestMetaInfo,
        override val attachments: List<Attachment> = emptyList()
    ) : Request, WithAttachments {
        override val role: Role = Role.User
    }

    /**
     * Represents a message generated by the assistant as a response.
     *
     * @property content The textual content of the assistant's response.
     * @property metaInfo Metadata related to the response, including token counts and timestamp.
     * @property attachments Optional media content attached to the message.
     * @property finishReason An optional explanation for why the assistant's response was finalized.
     * Defaults to null if not provided.
     * @property role The role associated with the response, which is fixed as `Role.Assistant`.
     */
    @Serializable
    public data class Assistant(
        override val content: String,
        override val metaInfo: ResponseMetaInfo,
        override val attachments: List<Attachment> = emptyList(),
        val finishReason: String? = null
    ) : Response, WithAttachments {
        override val role: Role = Role.Assistant

        override fun copy(updatedMetaInfo: ResponseMetaInfo): Assistant = this.copy(metaInfo = updatedMetaInfo)
    }

    /**
     * Represents messages exchanged with tools, either as calls or results.
     */
    @Serializable
    public sealed interface Tool : Message {
        /**
         * The unique identifier of the tool call.
         */
        public val id: String?

        /**
         * The name of the tool used.
         */
        public val tool: String

        /**
         * Represents a tool call message sent as a response.
         *
         * @property id The unique identifier of the tool call.
         * @property tool The name of the tool being called.
         * @property content The content of the tool call.
         * @property metaInfo Metadata related to the response, including token counts and timestamp.
         */
        @Serializable
        public data class Call(
            override val id: String?,
            override val tool: String,
            override val content: String,
            override val metaInfo: ResponseMetaInfo
        ) : Tool, Response {
            override val role: Role = Role.Tool

            /**
             * Lazily parses the content of the tool call as a JSON object.
             */
            val contentJson: JsonObject by lazy {
                Json.parseToJsonElement(content).jsonObject
            }

            override fun copy(updatedMetaInfo: ResponseMetaInfo): Call = this.copy(metaInfo = updatedMetaInfo)
        }

        /**
         * Represents the result of a tool call sent as a request.
         *
         * @property id The unique identifier of the tool result.
         * @property tool The name of the tool that provided the result.
         * @property content The content of the tool result.
         * @property metaInfo Metadata associated with the request, including timestamp information. Defaults to a new [RequestMetaInfo].
         */
        @Serializable
        public data class Result(
            override val id: String?,
            override val tool: String,
            override val content: String,
            override val metaInfo: RequestMetaInfo
        ) : Tool, Request {
            override val role: Role = Role.Tool
        }
    }

    /**
     * Represents a system-generated message.
     *
     * @property content The content of the system message.
     * @property metaInfo Metadata associated with the request, including timestamp information. Defaults to a new [RequestMetaInfo].
     *
     */
    @Serializable
    public data class System(
        override val content: String,
        override val metaInfo: RequestMetaInfo
    ) : Request {
        override val role: Role = Role.System
    }
}

/**
 * Meta-information associated with a message in a chat system.
 *
 * @property timestamp The timestamp [Instant] of when the message is created
 * since the Unix epoch. Defaults to the current system time.
 */
@Serializable
public sealed interface MessageMetaInfo {
    /**
     * Represents the timestamp of a message
     *
     * This property indicates the precise time when a message was created. It defaults
     * to the current system time if not explicitly set.
     */
    public val timestamp: Instant

    /**
     * Free-form information associated with a message.
     * Can be used to store custom metadata that doesn't fit into the standard fields.
     */
    public val metadata: JsonObject?
}

/**
 * Represents [MessageMetaInfo] specific to a request within the system.
 *
 * This class is an implementation of the [MessageMetaInfo] interface and provides
 * timestamp information for a request.
 *
 * @property timestamp The time at which the request metadata was created.
 * Defaults to the current system time if not provided.
 */
@Serializable
public data class RequestMetaInfo(
    override val timestamp: Instant,
    override val metadata: JsonObject? = null
) : MessageMetaInfo {
    /**
     * Companion object for `RequestMetaInfo` that provides factory methods and utilities related to creating instances.
     */
    public companion object {
        /**
         * Creates a RequestMetadata instance with a timestamp from the provided clock.
         *
         * @param clock The clock to use for generating the timestamp.
         * @return A new RequestMetadata instance with the timestamp from the provided clock.
         */
        public fun create(clock: Clock): RequestMetaInfo = RequestMetaInfo(clock.now())

        /**
         * An empty instance of [RequestMetaInfo] with the timestamp set to a distant past.
         */
        public val Empty: RequestMetaInfo = RequestMetaInfo(Instant.DISTANT_PAST)
    }
}

/**
 * Represents metadata associated with a response message in a chat system.
 *
 * This class provides details about the response, including the count of tokens
 * used in the response and the timestamp of when the response was created.
 * It implements the `MessageMetadata` interface, inheriting the timestamp property.
 *
 *
 * Example:
 * - Message 1: "Hello" (3 tokens) → tokensCount = 3
 * - Message 2: "How are you?" (4 tokens) → tokensCount = 3 + 4 = 7
 * - Message 3: "I am fine, thank you." (6 tokens) → tokensCount = 7 + 6 = 13
 *
 * @property totalTokensCount The total number of tokens involved in the response, including both input and output tokens, or null if not available.
 * @property inputTokensCount The number of tokens used in the input, or null if not available.
 * @property outputTokensCount The number of tokens generated in the output, or null if not available.
 * @property additionalInfo Additional metadata as a map of string keys to string values.
 *                          This can be used to store custom metadata that doesn't fit into the standard fields.
 * @property timestamp The timestamp indicating when the response was created.
 * Defaults to the current system time if not explicitly set.
 */
@Serializable
public data class ResponseMetaInfo(
    public override val timestamp: Instant,
    public val totalTokensCount: Int? = null,
    public val inputTokensCount: Int? = null,
    public val outputTokensCount: Int? = null,
    @Deprecated(
        "additionalInfo is deprecated, use metadata instead",
        ReplaceWith("metadata")
    )
    public val additionalInfo: Map<String, String> = emptyMap(),
    override val metadata: JsonObject? = null,
) : MessageMetaInfo {
    /**
     * Companion object for the ResponseMetaInfo class.
     * Provides utility methods to create ResponseMetaInfo instances.
     */
    public companion object {
        /**
         * Creates a ResponseMetadata instance with a timestamp from the provided clock.
         *
         * @param clock The clock to use for generating the timestamp.
         * @param totalTokensCount The total number of tokens involved in the response, including both input and output tokens.
         * @param inputTokensCount The number of tokens used in the input.
         * @param outputTokensCount The number of tokens generated in the output.
         * @param additionalInfo Deprecated: use [metadata] instead. Additional metadata as a map of string keys to string values.
         * @param metadata Additional metadata as a JSON object.
         * @return A new ResponseMetadata instance with the timestamp from the provided clock.
         */
        public fun create(
            clock: Clock,
            totalTokensCount: Int? = null,
            inputTokensCount: Int? = null,
            outputTokensCount: Int? = null,
            additionalInfo: Map<String, String> = emptyMap(),
            metadata: JsonObject? = null,
        ): ResponseMetaInfo =
            ResponseMetaInfo(clock.now(), totalTokensCount, inputTokensCount, outputTokensCount, additionalInfo, metadata)

        /**
         * An empty instance of the [ResponseMetaInfo] with the timestamp set to a distant past.
         */
        public val Empty: ResponseMetaInfo = ResponseMetaInfo(Instant.DISTANT_PAST)
    }
}
